import cv2
import mediapipe as mp
import time

class PoseDetector():
    def __init__(self,
                 static_image_mode=False,
                 model_complexity=1,
                 smooth_landmarks=True,
                 enable_segmentation=False,
                 smooth_segmentation=True,
                 min_detection_confidence=0.5,
                 min_tracking_confidence=0.5):
        self.static_image_mode = static_image_mode
        self.model_complexity = model_complexity
        self.smooth_landmarks = smooth_landmarks
        self.enable_segmentation = enable_segmentation
        self.smooth_segmentation = smooth_segmentation
        self.min_detection_confidence = min_detection_confidence
        self.min_tracking_confidence = min_tracking_confidence

        self.mpPose = mp.solutions.pose
        self.pose = self.mpPose.Pose(self.static_image_mode, self.smooth_landmarks,
                                     self.enable_segmentation, self.enable_segmentation,
                                     self.smooth_segmentation, self.min_detection_confidence, self.min_tracking_confidence,)
        self.mpDraw = mp.solutions.drawing_utils


    def findPose(self, frame, draw=True):

        frameRGB = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        self.results = self.pose.process(frameRGB)
        if self.results.pose_landmarks:
            if draw:
                self.mpDraw.draw_landmarks(frame, self.results.pose_landmarks, self.mpPose.POSE_CONNECTIONS)

        return frame


    def findPosition(self, frame, draw=True):
        """找点,0-32个点位，表示人体不同位置"""
        lists = []
        if self.results.pose_landmarks:
            for id, lm in enumerate(self.results.pose_landmarks.landmark):
                h, w, c = frame.shape
                cx, cy = int(lm.x * w), int(lm.y * h)
                lists.append([id, cx, cy])
                if draw:
                    cv2.circle(frame, (cx, cy), 5, (255, 0, 0), cv2.FILLED)
        return lists

def main():
    cap = cv2.VideoCapture(0)
    Ptime = 0
    detector = PoseDetector()
    while True:
        ret, frame = cap.read()
        frame = detector.findPose(frame)
        lists = detector.findPosition(frame,draw=False)
        if len(lists) != 0:
            print(lists[0])
            cv2.circle(frame, (lists[0][1], lists[0][2]), 5, (255, 0, 0), cv2.FILLED)

        Ctime = time.time()
        fps = 1 / (Ctime - Ptime)
        Ptime = Ctime
        cv2.putText(frame, 'FPS:' + str(int(fps)), (10, 70), cv2.FONT_HERSHEY_PLAIN, 3, (0, 255, 0), 3)
        cv2.imshow('frame', frame)
        if cv2.waitKey(1) == ord('q'):
            break
    cap.release()
    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()
